#!/bin/env python
import sys
import os
from git import Repo

def gen_python_cuda_base_file(base_imag="update_cuda11.8_ubuntu22.04",python_version="3.10"):
    pyquda_def=f"""Bootstrap:localimage
From: {base_imag}.sif
%setup
%files
%post
    pip install -U setuptools
    pip install -U pip
    pip install numpy scipy sympy matplotlib opt_einsum xarray h5py h5netcdf gmpy2
    pip install psutil gvar lsqfit corrfitter
    pip install mpi4py 
%environment
    export LD_LIBRARY_PATH=$CUDA_PATH/lib64:$LD_LIBRARY_PATH
    export CUDA_PATH=/usr/local/cuda:$CUDA_PATH
%runscript
    """
    pyquda_imag = f"python{python_version}_" + base_imag
    with open(pyquda_imag + ".def","w") as f:
        f.writelines(pyquda_def)
    

def gen_quda_define_files(quda_commit=None,base_imag="python3.10_update_cuda11.8_ubuntu22.04",
                                        sm="sm_35",):
    if sm == "sm_35":
        cupy_def = """
    cd /src/cupy
    git config --global --add safe.directory /src/cupy
    git checkout v13
    cp /src/_routines_linalg.pyx /src/cupy/cupy/_core/
    pip install "cython<3"
    export CUDA_PATH=/usr/local/cuda:$CUDA_PATH
    export LD_LIBRARY_PATH=$CUDA_PATH/lib64:$LD_LIBRARY_PATH
    pip install .
"""
    else:
        cupy_def="""
    pip install cython
    pip install "cupy-cuda11x>=12"
"""
    pyquda_def=f"""Bootstrap:localimage
From: {base_imag}.sif
%setup

%files
    /home/guilongcheng/source_codes/cupy /src/
    /home/guilongcheng/lattice_cloud/source_codes/*.* /src/
    /home/guilongcheng/lattice_cloud/source_codes/quda /src/
%post
    {cupy_def}
    cd /src/
    ./build_pyquda.sh {sm} 
    export QUDA_PATH=/usr/local/quda
%environment
    export QUDA_PATH=/usr/local/quda
%runscript
    """
    pyquda_imag = "quda_"+ f"{quda_commit}_" + sm + "_" + base_imag
    with open(pyquda_imag + ".def","w") as f:
        f.writelines(pyquda_def)


gen_python_cuda_base_file()
# 初始化仓库对象，这里替换成你的本地仓库路径
repo_path = os.path.expanduser("~/lattice_cloud/source_codes/quda")
repo = Repo(repo_path)
# 获取最近的commit
latest_commit = repo.head.commit
# 获取短哈希
short_hash = latest_commit.hexsha[:7]
print(short_hash)
for i in ["sm_35","sm_60","sm_70","sm_80"]:
    gen_quda_define_files(quda_commit=short_hash,sm=i)    
